
package com.shxs.LLMConnect.chatGLM;
import com.shxs.LLMConnect.LLMConnector;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.ChatCompletionRequest;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ChatMessageRole;
import com.zhipu.oapi.service.v4.model.ModelApiResponse;

import java.util.ArrayList;
import java.util.List;

public class ChatGLMConnector implements LLMConnector {
    ClientV4 client;

    String modelType = Constants.ModelChatGLM4;

    boolean isStream = false;

    float temperature = 0.5f;

    int maxTokens = 8192;


    public ChatGLMConnector(String apikey) {
        client = new ClientV4.Builder(apikey).build();
    }

    public ChatGLMConnector ModelName(String model) {
        modelType = model;
        return this;
    }

    public ChatGLMConnector isStream(boolean isStream) {
        this.isStream = isStream;
        return this;
    }

    public ChatGLMConnector temperature(float temperature) {
        if (temperature < 1.0f && temperature > 0.0f)
            this.temperature = temperature;
        return this;
    }

    public ChatGLMConnector maxTokens(int maxTokens) {
        this.maxTokens = maxTokens;
        return this;
    }

    @Override
    public String askWithHistory(String line) {
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), line);
        messages.add(chatMessage);
        line=ask(messages.toString());
        chatMessage = new ChatMessage(ChatMessageRole.ASSISTANT.value(), line);
        messages.add(chatMessage);
        return line;
    }

    @Override
    public String ask(String message) {

        String requestId = String.format("Novel2Sound-%d", System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(modelType)
                .stream(isStream)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .temperature(temperature)
                .maxTokens(maxTokens)
                .build();
        ModelApiResponse invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        try {
            JSONObject js = JSON.parseObject(JSON.toJSONString(invokeModelApiResp));
            String line = js.getString("success");
            if (!line.equals("true")) return null;
            line = js.getJSONObject("data")
                    .getJSONArray("choices").getJSONObject(0)
                    .getJSONObject("message")
                    .getString("content");
            return line;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    List<ChatMessage> messages = new ArrayList<>();
    @Override
    public int getHistoryLength() {
        return messages.size();
    }

}
